#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
增量图片上传脚本
高性能多线程上传图片到云服务器，支持断点续传和增量更新
"""

import os
import sys
import json
import hashlib
import logging
import threading
import time
from pathlib import Path
from typing import Dict, List, Optional, Tuple
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
import paramiko
from paramiko import SSHClient, SFTPClient
from tqdm import tqdm


@dataclass
class UploadConfig:
    """上传配置"""
    # 服务器连接配置
    hostname: str
    port: int = 22
    username: str
    password: Optional[str] = None
    private_key_path: Optional[str] = None
    
    # 路径配置
    local_base_path: str
    remote_base_path: str
    
    # 上传配置
    max_workers: int = 15  # 并发线程数
    chunk_size: int = 32768  # 传输块大小
    retry_times: int = 3  # 重试次数
    retry_delay: int = 1  # 重试延迟（秒）
    
    # 增量配置
    check_method: str = "size"  # size, md5, both
    skip_existing: bool = True  # 跳过已存在的相同文件


@dataclass
class UploadStats:
    """上传统计信息"""
    total_files: int = 0
    uploaded_files: int = 0
    skipped_files: int = 0
    failed_files: int = 0
    total_size: int = 0
    uploaded_size: int = 0
    start_time: float = 0
    
    def get_speed_mbps(self) -> float:
        """获取上传速度（MB/s）"""
        elapsed = time.time() - self.start_time
        if elapsed > 0:
            return (self.uploaded_size / 1024 / 1024) / elapsed
        return 0


class IncrementalImageUploader:
    """增量图片上传器"""
    
    def __init__(self, config: UploadConfig):
        self.config = config
        self.stats = UploadStats()
        self.progress_lock = threading.Lock()
        self.progress_bar = None
        
        # 设置日志
        self._setup_logging()
        
        # 状态文件
        self.state_file = self._get_state_file_path()
        self.uploaded_files = self._load_upload_state()
        
        # SSH连接池
        self.ssh_pool = []
        self.pool_lock = threading.Lock()
    
    def _setup_logging(self):
        """设置日志记录"""
        script_dir = Path(__file__).parent
        project_dir = script_dir.parent
        log_dir = project_dir / "logs"
        log_dir.mkdir(exist_ok=True)
        
        log_file = log_dir / f"image_upload_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log"
        
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(threadName)s - %(levelname)s - %(message)s',
            handlers=[
                logging.FileHandler(log_file, encoding='utf-8'),
                logging.StreamHandler()
            ]
        )
        self.logger = logging.getLogger(__name__)
    
    def _get_state_file_path(self) -> Path:
        """获取状态文件路径"""
        script_dir = Path(__file__).parent
        project_dir = script_dir.parent
        state_dir = project_dir / "upload_state"
        state_dir.mkdir(exist_ok=True)
        
        # 基于配置生成唯一的状态文件名
        config_hash = hashlib.md5(
            f"{self.config.hostname}_{self.config.remote_base_path}".encode()
        ).hexdigest()[:8]
        
        return state_dir / f"upload_state_{config_hash}.json"
    
    def _load_upload_state(self) -> Dict[str, Dict]:
        """加载上传状态"""
        if self.state_file.exists():
            try:
                with open(self.state_file, 'r', encoding='utf-8') as f:
                    return json.load(f)
            except Exception as e:
                self.logger.warning(f"加载状态文件失败: {e}")
        return {}
    
    def _save_upload_state(self):
        """保存上传状态"""
        try:
            with open(self.state_file, 'w', encoding='utf-8') as f:
                json.dump(self.uploaded_files, f, ensure_ascii=False, indent=2)
        except Exception as e:
            self.logger.error(f"保存状态文件失败: {e}")
    
    def _create_ssh_connection(self) -> SSHClient:
        """创建SSH连接"""
        ssh = SSHClient()
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        
        try:
            if self.config.private_key_path:
                # 使用私钥认证
                key = paramiko.RSAKey.from_private_key_file(self.config.private_key_path)
                ssh.connect(
                    hostname=self.config.hostname,
                    port=self.config.port,
                    username=self.config.username,
                    pkey=key,
                    timeout=30
                )
            else:
                # 使用密码认证
                ssh.connect(
                    hostname=self.config.hostname,
                    port=self.config.port,
                    username=self.config.username,
                    password=self.config.password,
                    timeout=30
                )
            
            return ssh
        except Exception as e:
            self.logger.error(f"SSH连接失败: {e}")
            raise
    
    def _get_ssh_connection(self) -> SSHClient:
        """从连接池获取SSH连接"""
        with self.pool_lock:
            if self.ssh_pool:
                return self.ssh_pool.pop()
        
        # 连接池为空，创建新连接
        return self._create_ssh_connection()
    
    def _return_ssh_connection(self, ssh: SSHClient):
        """归还SSH连接到连接池"""
        with self.pool_lock:
            if len(self.ssh_pool) < self.config.max_workers:
                self.ssh_pool.append(ssh)
            else:
                try:
                    ssh.close()
                except:
                    pass
    
    def _calculate_file_md5(self, file_path: Path) -> str:
        """计算文件MD5"""
        hash_md5 = hashlib.md5()
        with open(file_path, "rb") as f:
            for chunk in iter(lambda: f.read(4096), b""):
                hash_md5.update(chunk)
        return hash_md5.hexdigest()
    
    def _get_remote_file_info(self, sftp: SFTPClient, remote_path: str) -> Optional[Dict]:
        """获取远程文件信息"""
        try:
            stat = sftp.stat(remote_path)
            return {
                'size': stat.st_size,
                'mtime': stat.st_mtime
            }
        except FileNotFoundError:
            return None
        except Exception as e:
            self.logger.warning(f"获取远程文件信息失败 {remote_path}: {e}")
            return None
    
    def _should_upload_file(self, local_file: Path, remote_path: str, sftp: SFTPClient) -> bool:
        """判断是否需要上传文件"""
        local_size = local_file.stat().st_size
        relative_path = str(local_file.relative_to(self.config.local_base_path))
        
        # 检查本地状态记录
        if relative_path in self.uploaded_files:
            uploaded_info = self.uploaded_files[relative_path]
            if uploaded_info.get('size') == local_size:
                if self.config.check_method == 'size':
                    return False
                elif self.config.check_method == 'md5':
                    local_md5 = self._calculate_file_md5(local_file)
                    if uploaded_info.get('md5') == local_md5:
                        return False
        
        # 检查远程文件
        if self.config.skip_existing:
            remote_info = self._get_remote_file_info(sftp, remote_path)
            if remote_info:
                if self.config.check_method == 'size' and remote_info['size'] == local_size:
                    # 记录到状态文件
                    self.uploaded_files[relative_path] = {
                        'size': local_size,
                        'upload_time': time.time()
                    }
                    return False
                elif self.config.check_method in ['md5', 'both']:
                    # MD5检查需要下载文件或使用其他方法，这里简化处理
                    if remote_info['size'] == local_size:
                        self.uploaded_files[relative_path] = {
                            'size': local_size,
                            'upload_time': time.time()
                        }
                        return False
        
        return True
    
    def _upload_single_file(self, local_file: Path, remote_path: str) -> bool:
        """上传单个文件"""
        ssh = None
        try:
            ssh = self._get_ssh_connection()
            sftp = ssh.open_sftp()
            
            # 检查是否需要上传
            if not self._should_upload_file(local_file, remote_path, sftp):
                with self.progress_lock:
                    self.stats.skipped_files += 1
                    if self.progress_bar:
                        self.progress_bar.set_postfix({
                            'uploaded': self.stats.uploaded_files,
                            'skipped': self.stats.skipped_files,
                            'failed': self.stats.failed_files,
                            'speed': f"{self.stats.get_speed_mbps():.1f}MB/s"
                        })
                        self.progress_bar.update(1)
                return True
            
            # 创建远程目录
            remote_dir = os.path.dirname(remote_path)
            try:
                sftp.makedirs(remote_dir)
            except:
                pass  # 目录可能已存在
            
            # 上传文件
            file_size = local_file.stat().st_size
            
            def progress_callback(transferred, total):
                pass  # 可以添加单文件进度回调
            
            sftp.put(str(local_file), remote_path, callback=progress_callback)
            
            # 更新统计和状态
            relative_path = str(local_file.relative_to(self.config.local_base_path))
            upload_info = {
                'size': file_size,
                'upload_time': time.time()
            }
            
            if self.config.check_method in ['md5', 'both']:
                upload_info['md5'] = self._calculate_file_md5(local_file)
            
            self.uploaded_files[relative_path] = upload_info
            
            with self.progress_lock:
                self.stats.uploaded_files += 1
                self.stats.uploaded_size += file_size
                
                if self.progress_bar:
                    self.progress_bar.set_postfix({
                        'uploaded': self.stats.uploaded_files,
                        'skipped': self.stats.skipped_files,
                        'failed': self.stats.failed_files,
                        'speed': f"{self.stats.get_speed_mbps():.1f}MB/s"
                    })
                    self.progress_bar.update(1)
            
            sftp.close()
            return True
            
        except Exception as e:
            self.logger.error(f"上传文件失败 {local_file}: {e}")
            with self.progress_lock:
                self.stats.failed_files += 1
                if self.progress_bar:
                    self.progress_bar.set_postfix({
                        'uploaded': self.stats.uploaded_files,
                        'skipped': self.stats.skipped_files,
                        'failed': self.stats.failed_files,
                        'speed': f"{self.stats.get_speed_mbps():.1f}MB/s"
                    })
                    self.progress_bar.update(1)
            return False
        finally:
            if ssh:
                self._return_ssh_connection(ssh)
    
    def _scan_local_files(self) -> List[Tuple[Path, str]]:
        """扫描本地文件"""
        self.logger.info("扫描本地文件...")
        
        local_path = Path(self.config.local_base_path)
        files_to_upload = []
        
        # 支持的图片格式
        image_extensions = {'.jpg', '.jpeg', '.png', '.gif', '.bmp', '.webp', '.tiff', '.tif'}
        
        for file_path in local_path.rglob('*'):
            if file_path.is_file() and file_path.suffix.lower() in image_extensions:
                # 计算相对路径
                relative_path = file_path.relative_to(local_path)
                remote_path = os.path.join(self.config.remote_base_path, str(relative_path)).replace('\\', '/')
                
                files_to_upload.append((file_path, remote_path))
                self.stats.total_size += file_path.stat().st_size
        
        self.stats.total_files = len(files_to_upload)
        self.logger.info(f"找到 {self.stats.total_files} 个图片文件，总大小 {self.stats.total_size/1024/1024/1024:.2f}GB")
        
        return files_to_upload
    
    def upload(self) -> bool:
        """执行上传"""
        self.logger.info("=" * 60)
        self.logger.info("开始增量图片上传")
        self.logger.info(f"本地路径: {self.config.local_base_path}")
        self.logger.info(f"远程路径: {self.config.hostname}:{self.config.remote_base_path}")
        self.logger.info(f"并发线程: {self.config.max_workers}")
        self.logger.info("=" * 60)
        
        # 扫描文件
        files_to_upload = self._scan_local_files()
        if not files_to_upload:
            self.logger.info("没有找到需要上传的文件")
            return True
        
        # 开始上传
        self.stats.start_time = time.time()
        
        with tqdm(total=self.stats.total_files, desc="上传进度", unit="file") as pbar:
            self.progress_bar = pbar
            
            with ThreadPoolExecutor(max_workers=self.config.max_workers) as executor:
                # 提交所有上传任务
                future_to_file = {
                    executor.submit(self._upload_single_file, local_file, remote_path): (local_file, remote_path)
                    for local_file, remote_path in files_to_upload
                }
                
                # 处理完成的任务
                for future in as_completed(future_to_file):
                    local_file, remote_path = future_to_file[future]
                    try:
                        success = future.result()
                        if success:
                            # 定期保存状态
                            if (self.stats.uploaded_files + self.stats.skipped_files) % 100 == 0:
                                self._save_upload_state()
                    except Exception as e:
                        self.logger.error(f"处理文件时出错 {local_file}: {e}")
        
        # 清理连接池
        with self.pool_lock:
            for ssh in self.ssh_pool:
                try:
                    ssh.close()
                except:
                    pass
            self.ssh_pool.clear()
        
        # 保存最终状态
        self._save_upload_state()
        
        # 输出统计信息
        elapsed_time = time.time() - self.stats.start_time
        self.logger.info("=" * 60)
        self.logger.info("上传完成！")
        self.logger.info(f"总文件数: {self.stats.total_files}")
        self.logger.info(f"上传文件: {self.stats.uploaded_files}")
        self.logger.info(f"跳过文件: {self.stats.skipped_files}")
        self.logger.info(f"失败文件: {self.stats.failed_files}")
        self.logger.info(f"上传大小: {self.stats.uploaded_size/1024/1024/1024:.2f}GB")
        self.logger.info(f"用时: {elapsed_time/60:.1f}分钟")
        self.logger.info(f"平均速度: {self.stats.get_speed_mbps():.1f}MB/s")
        self.logger.info("=" * 60)
        
        return self.stats.failed_files == 0


def create_config_template():
    """创建配置文件模板"""
    config_template = {
        "hostname": "your-server.com",
        "port": 22,
        "username": "your-username",
        "password": "your-password",
        "private_key_path": null,
        "local_base_path": "F:/0109_repository/gitee_repository/chuhai/chuhai_scraper/storage/vehicles/picknbuy24",
        "remote_base_path": "/var/www/images/vehicles/picknbuy24",
        "max_workers": 15,
        "chunk_size": 32768,
        "retry_times": 3,
        "retry_delay": 1,
        "check_method": "size",
        "skip_existing": true
    }
    
    config_file = Path("chuhai_scraper/config/upload_config.json")
    config_file.parent.mkdir(exist_ok=True)
    
    with open(config_file, 'w', encoding='utf-8') as f:
        json.dump(config_template, f, ensure_ascii=False, indent=2)
    
    print(f"配置文件模板已创建: {config_file}")
    print("请编辑配置文件后重新运行脚本")


def main():
    """主函数"""
    print("高性能增量图片上传脚本")
    print("=" * 60)
    
    config_file = Path("chuhai_scraper/config/upload_config.json")
    
    # 检查配置文件
    if not config_file.exists():
        print("配置文件不存在，正在创建模板...")
        create_config_template()
        return
    
    # 加载配置
    try:
        with open(config_file, 'r', encoding='utf-8') as f:
            config_data = json.load(f)
        
        config = UploadConfig(**config_data)
    except Exception as e:
        print(f"加载配置文件失败: {e}")
        return
    
    # 验证配置
    if config.hostname == "your-server.com":
        print("请先编辑配置文件 chuhai_scraper/config/upload_config.json")
        return
    
    # 确认上传
    print(f"本地路径: {config.local_base_path}")
    print(f"远程服务器: {config.hostname}:{config.remote_base_path}")
    print(f"并发线程: {config.max_workers}")
    print(f"检查方式: {config.check_method}")
    
    response = input("\n确认开始上传？(y/N): ").strip().lower()
    if response not in ['y', 'yes']:
        print("上传已取消")
        return
    
    # 执行上传
    uploader = IncrementalImageUploader(config)
    success = uploader.upload()
    
    if success:
        print("\n✅ 上传任务成功完成！")
    else:
        print("\n❌ 上传过程中出现错误，请查看日志文件")


if __name__ == "__main__":
    main()
